from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from patagent.llm import chat_llm
from patagent.prompt import ANSWER_GRADE_PROMPT

class GradeAnswer(BaseModel):
    """Binary score to assess answer addresses question."""

    binary_score: str = Field(
        description="Answer addresses the question, 'yes' or 'no'"
    )

# LLM with function call
structured_llm_grader = chat_llm.with_structured_output(GradeAnswer)

answer_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", ANSWER_GRADE_PROMPT),
        ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
    ]
)

answer_grader = answer_prompt | structured_llm_grader
# answer_grader.invoke({"question": question, "generation": generation})